#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from retrying import retry
import requests
import sys
import os


@retry
def EfetchFasta(ids, parentpath, prefix, index, count):
    url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi'
    payload = {"db": "nuccore", "id": ids, "rettype": "fasta"}
    r = requests.post(url, stream=True, data=payload)
    if r.status_code == 200:
        with open(parentpath + "/" + prefix + "_" + index + ".fasta", 'wb') as fd:
            for chunk in r.iter_content(1024 * 1024):
                fd.write(chunk)
            fd.close()
        if not Verify_Fasta(parentpath + "/" + prefix + "_" + index + ".fasta", count):
            print("Invalid fasta. Try to redownload ...")
            raise IOError("Invalid fasta.")
        else:
            print(prefix + "_" + index + ".fasta downloaded.")
    else:
        print("Web Error. Retrying ...")
        raise IOError("Web Error.")


def Verify_Fasta(filename, count):
    import subprocess

    cmd2 = "grep '>' " + filename + " | wc -l"
    p2 = subprocess.Popen(cmd2, stdout=subprocess.PIPE, shell=True)
    p2.wait()
    result2 = ""
    for line in p2.stdout.readlines():
        line = line.decode("utf-8").strip()
        result2 += line
    if int(result2.strip()) == count:
        cmd = "grep '<html>' " + filename
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
        p.wait()
        result = ""
        for line in p.stdout.readlines():
            line = line.decode("utf-8").rstrip()
            result += line
        if result.find("<html>") != -1:
            print(filename + " invalid because of <html>")
            return False
        elif result.find("No such file or directory") != -1:
            print(filename + " invalid because of no file")
            return False
        else:
            return True
    else:
        print(filename + " invalid because not enough fasta, real fasta count:" + result2.strip())
        return False


if __name__ == '__main__':
    parentpath = os.path.abspath(sys.argv[1])[:os.path.abspath(sys.argv[1]).rindex("/")]
    basename = os.path.basename(sys.argv[1])
    os.makedirs(parentpath + "/fasta", exist_ok=True)
    prefix = basename.split(".")[0]

    infile = open(sys.argv[1], "r")

    index = 0
    count = 0
    ids = ""

    for line in infile:
        line = line.rstrip("\n")
        fasta_id = line
        ids += fasta_id + ","
        count += 1
        if count == 5000:
            if os.path.exists(parentpath + "/fasta/" + prefix + "_" + str(index).zfill(5) + ".fasta"):
                print(prefix + "_" + str(index).zfill(5) + ".fasta exists, skip.")
            else:
                EfetchFasta(ids[:-1], parentpath + "/fasta", prefix, str(index).zfill(5), count)
            count = 0
            ids = ""
            index += 1
    if count > 0:
        if os.path.exists(parentpath + "/fasta/" + prefix + "_" + str(index).zfill(5) + ".fasta"):
            print(prefix + "_" + str(index).zfill(5) + ".fasta exists, skip.")
        else:
            EfetchFasta(ids[:-1], parentpath + "/fasta", prefix, str(index).zfill(5), count)
    infile.close()
